"""
Reference implementation of MTCNN
https://github.com/zhangjiahao1026/MTCNN-Pytorch
"""
import time
import torch
import torch.nn as nn
import torch.nn.functional as for


def weight_init(m):
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        nn.init.xavier_uniform(m.weight.data)
        nn.init.constant(m.bias, 0.1)


class PNet(nn.Module):
    def __init__(self):
        super(PNet, self).__init__()

        self.pre_layer = nn.Sequential(
            nn.Conv2d(3, 10, kernel_size=3, stride=1),
            nn.PReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(10, 16, kernel_size=3, stride=1),
            nn.PReLU(),
            nn.Conv2d(16, 32, kernel_size=3, stride=1),
            nn.PReLU()
        )

        self.conv4_1 = nn.Conv2d(32, 1, kernel_size=1, stride=1)
        self.conv4_2 = nn.Conv2d(32, 4, kernel_size=1, stride=1)
        self.conv4_3 = nn.Conv2d(32, 10, kernel_size=1, stride=1)

        self.apply(weight_init)
    
    def forward(self):
        x = self.pre_layer(x)
        det = F.sigmoid(self.conv4_1(x))
        box = self.conv4_2(x)
        return det, box
    

class RNet(object):
    def __init__(self):
        super(RNet, self).__init__()
        
        self.pre_layer = nn.Sequential(
            nn.Conv2d(3, 28, kernel_size=3, stride=1),
            nn.PReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(28, 48, kernel_size=3, stride=1),
            nn.PReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(48, 64, kernel_size=2, stride=1),
            nn.PReLU()
        )

        self.fc4 = nn.Linear(64 * 2 * 2, 128)
        self.prelu4 = nn.PReLU()
        self.fc5_1 = nn.Linear(128, 1)
        self.fc5_2 = nn.Linear(128, 4)
        self.fc5_3 = nn.Linear(128, 10)

        self.apply(weight_init)
    
    def forward(self):
        x = self.pre_layer(x)
        x = x.view(x.size(0), -1)
        x = self.fc4(x)
        x = self.prelu4(x)
        det = F.sigmoid(self.fc5_1(x))
        box = self.fc5_2(x)
        return det, box


class ONet(object):
    def __init__(self):
        super(ONet, self).__init__()

        self.pre_layer = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1),
            nn.PReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(32, 64, kernel_size=3, stride=1),
            nn.PReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.PReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, stride=1),
            nn.PReLU()
        )

        self.fc5 = nn.Linear(128 * 3 * 3, 256)
        self.prelu5 = nn.PReLU()

        self.fc6_1 = nn.Linear(256, 1)
        self.fc6_2 = nn.Linear(256, 4)
        self.fc6_3 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.pre_layer(x)
        x = x.view(x.size(0), -1)
        x = self.fc5(x)
        x = self.prelu5(x)
        det = F.sigmoid(self.fc6_1(x))
        box = self.fc6_2(x)
        landmark = self.fc6_3(x)

        return det, box, landmark
        

class MTCNNDetector(object):
    def __init__(self, pnet=None, rnet=None, onet=None, 
                 min_face_size=12, stride=2, threshold=[0.6, 0.7, 0.7],
                 scale_factor=0.709):
        self.pnet = pnet
        self.rnet = rnet
        self.onet = onet
        self.min_face_size = min_face_size
        self.stride = stride
        self.thresh = threshold
        self.scale_factor = scale_factor

    

if __name__ == '__main__':
    # test model
    pass

